"""
Demo script showing how to use the Attribution Agent.
"""
import os
import sys
sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

from attribution_agent import AttributionAgent
from examples.sample_data import generate_ecommerce_data, generate_saas_data
from datetime import datetime, timedelta
import pandas as pd


def demo_ecommerce_analysis():
    """Demo: E-commerce revenue attribution analysis."""
    print("=" * 80)
    print("DEMO 1: E-commerce Revenue Attribution")
    print("=" * 80)

    # Generate sample data
    print("\n1. Generating sample e-commerce data...")
    data = generate_ecommerce_data(start_date="2025-09-01", days=60)
    print(f"   Generated {len(data)} rows of data")
    print(f"   Date range: {data['date'].min()} to {data['date'].max()}")
    print(f"   Dimensions: {[col for col in data.columns if data[col].dtype == 'object' and col != 'date']}")

    # Initialize agent
    print("\n2. Initializing Attribution Agent...")
    agent = AttributionAgent(config={
        'llm_model': 'openai/gpt-4o',  # OpenRouter format (or use 'gpt-4' for direct OpenAI)
        'max_drill_depth': 2,
        'anomaly_threshold': 2.0,
        'enable_llm_reasoning': True  # Set to False if no API key
    })

    # Run analysis
    print("\n3. Running attribution analysis...")
    query = """
    分析最近两周的收入表现，并与前两周进行比较。
    识别导致收入变化的原因。
    """

    result = agent.analyze(query=query, data=data)

    # Display results
    print("\n4. Analysis Results:")
    print("-" * 80)

    # Anomalies
    anomalies = result.get('anomalies', [])
    print(f"\n   Anomalies Detected: {len(anomalies)}")
    for i, anomaly in enumerate(anomalies[:3], 1):
        print(f"   {i}. {anomaly.get('timestamp', 'N/A')}: "
              f"{anomaly.get('deviation_percentage', 0):.1f}% deviation "
              f"({anomaly.get('severity', 'unknown')} severity)")

    # Drill-down path
    drill_path = result.get('drill_down_path', [])
    print(f"\n   Drill-Down Levels: {len(drill_path)}")
    for level, drill in enumerate(drill_path):
        print(f"\n   Level {level} - {drill.get('dimension', 'unknown')}:")
        for contrib in drill.get('top_contributors', [])[:3]:
            print(f"      • {contrib.get('value', 'N/A')}: "
                  f"{contrib.get('contribution_percentage', 0):.1f}% contribution")

    # Root causes
    root_causes = result.get('root_causes', [])
    print(f"\n   Root Causes: {len(root_causes)}")
    for i, cause in enumerate(root_causes, 1):
        print(f"   {i}. {cause.get('cause', 'Unknown')}")
        print(f"      Evidence: {cause.get('evidence', 'N/A')}")
        print(f"      Impact: {cause.get('impact', 'unknown')}")

    # Final report
    print("\n5. Final Report:")
    print("-" * 80)
    report = agent.get_report(result)
    print(report)

    return result


def demo_saas_mrr_analysis():
    """Demo: SaaS MRR attribution analysis."""
    print("\n\n" + "=" * 80)
    print("DEMO 2: SaaS MRR Attribution")
    print("=" * 80)

    # Generate sample data
    print("\n1. Generating sample SaaS data...")
    data = generate_saas_data(days=60)
    print(f"   Generated {len(data)} rows of data")

    # Initialize agent (without LLM for faster demo)
    print("\n2. Initializing Attribution Agent (no LLM)...")
    agent = AttributionAgent(config={
        'max_drill_depth': 2,
        'enable_llm_reasoning': False  # Faster, no API required
    })

    # Run analysis
    print("\n3. Running attribution analysis...")
    query = "Analyze MRR changes in the last 2 weeks"

    result = agent.analyze(query=query, data=data)

    # Display summary
    print("\n4. Analysis Summary:")
    print("-" * 80)

    drill_path = result.get('drill_down_path', [])
    if drill_path:
        first_level = drill_path[0]
        print(f"\n   Top dimension: {first_level.get('dimension', 'unknown')}")
        print("   Top contributors:")
        for contrib in first_level.get('top_contributors', [])[:5]:
            print(f"      • {contrib.get('value', 'N/A')}: "
                  f"{contrib.get('contribution_percentage', 0):.1f}%")

    # Report
    print("\n5. Report:")
    print("-" * 80)
    print(agent.get_report(result))

    return result


def demo_streaming():
    """Demo: Streaming analysis to see progress."""
    print("\n\n" + "=" * 80)
    print("DEMO 3: Streaming Analysis")
    print("=" * 80)

    data = generate_ecommerce_data(days=30)

    agent = AttributionAgent(config={'enable_llm_reasoning': False})

    print("\n Streaming analysis progress...\n")

    for i, state_update in enumerate(agent.stream_analyze(
        query="What's driving revenue changes?",
        data=data
    ), 1):
        node_name = list(state_update.keys())[0] if state_update else "unknown"
        print(f" Step {i}: {node_name}")

    print("\n Analysis complete!")


if __name__ == "__main__":
    print("Attribution Agent Demo\n")
    print("NOTE: Set OPENROUTER_API_KEY or OPENAI_API_KEY to enable LLM reasoning")
    print("      Or set enable_llm_reasoning=False for faster demo without API")
    print("      Get OpenRouter API key at: https://openrouter.ai/\n")

    # Run demos
    try:
        # Demo 1: E-commerce with full LLM (if API key available)
        has_api_key = os.getenv('OPENROUTER_API_KEY') or os.getenv('OPENAI_API_KEY')
        if has_api_key:
            demo_ecommerce_analysis()
        else:
            print("Skipping Demo 1 (requires OPENROUTER_API_KEY or OPENAI_API_KEY)")

        # Demo 2: SaaS without LLM (always runs)
        demo_saas_mrr_analysis()

        # Demo 3: Streaming
        demo_streaming()

    except Exception as e:
        print(f"\nError during demo: {str(e)}")
        import traceback
        traceback.print_exc()
